"""
Test if UPath works with Docling's save_as_json artifacts_dir parameter.

Usage:
1. Start MinIO: docker run -d \
    --name minio \
    -p 9000:9000 \
    -p 9001:9001 \
    -e MINIO_ROOT_USER=minioadmin \
    -e MINIO_ROOT_PASSWORD=minioadmin \
    minio/minio server /data --console-address ":9001"
2. Run: python test_upath_save.py
"""

import os
from pathlib import Path

from upath import UPath
from docling_core.types.doc import DoclingDocument, ImageRefMode

from docling.datamodel.base_models import InputFormat
from docling.datamodel.pipeline_options import PdfPipelineOptions
from docling.document_converter import DocumentConverter, PdfFormatOption


def test_type_compatibility():
    """Check if UPath passes isinstance checks."""
    local_path = Path("/tmp/test")
    s3_path = UPath("s3://bucket/test")

    print("=== Type Compatibility ===")
    print(f"Path instance check:     isinstance(Path(...), Path) = {isinstance(local_path, Path)}")
    print(f"UPath instance check:    isinstance(UPath(...), Path) = {isinstance(s3_path, Path)}")
    print(f"UPath has write_bytes:   {hasattr(s3_path, 'write_bytes')}")
    print(f"UPath has mkdir:         {hasattr(s3_path, 'mkdir')}")
    print(f"UPath has __truediv__:   {hasattr(s3_path, '__truediv__')}")
    print()


def convert_doc():
    """Convert a PDF with picture images enabled."""
    pipeline_options = PdfPipelineOptions()
    pipeline_options.generate_page_images = True
    pipeline_options.generate_picture_images = True

    converter = DocumentConverter(
        format_options={
            InputFormat.PDF: PdfFormatOption(pipeline_options=pipeline_options)
        }
    )

    result = converter.convert("../tests/data/pdf/2206.01062.pdf")
    return result.document


def test_local_path(doc: DoclingDocument):
    """Test save_as_json with regular Path."""
    print("=== Test: Local Path ===")
    output_dir = Path("scratch/local_test")
    output_dir.mkdir(parents=True, exist_ok=True)

    json_path = output_dir / "doc.json"
    artifacts_dir = output_dir / "images"

    try:
        doc.save_as_json(
            json_path,
            artifacts_dir=artifacts_dir,
            image_mode=ImageRefMode.REFERENCED
        )
        print(f"SUCCESS: Saved to {json_path}")
        print(f"  Images dir: {artifacts_dir}")
        if artifacts_dir.exists():
            images = list(artifacts_dir.glob("*.png"))
            print(f"  Image count: {len(images)}")
    except Exception as e:
        print(f"FAILED: {e}")
    print()


def test_upath_local(doc: DoclingDocument):
    """Test save_as_json with UPath (local filesystem)."""
    print("=== Test: UPath Local ===")
    output_dir = UPath("scratch/upath_local_test")
    output_dir.mkdir(parents=True, exist_ok=True)

    json_path = output_dir / "doc.json"
    artifacts_dir = output_dir / "images"

    try:
        doc.save_as_json(
            json_path,
            artifacts_dir=artifacts_dir,
            image_mode=ImageRefMode.REFERENCED
        )
        print(f"SUCCESS: Saved to {json_path}")
        print(f"  Images dir: {artifacts_dir}")
        if artifacts_dir.exists():
            images = list(artifacts_dir.glob("*.png"))
            print(f"  Image count: {len(images)}")
    except TypeError as e:
        print(f"TYPE ERROR (expected if Path type check fails): {e}")
    except Exception as e:
        print(f"FAILED: {type(e).__name__}: {e}")
    print()


def test_upath_s3(doc: DoclingDocument):
    """Test save_as_json with UPath pointing to MinIO/S3."""
    print("=== Test: UPath S3 (MinIO) ===")

    # Configure for local MinIO
    os.environ["AWS_ACCESS_KEY_ID"] = "minioadmin"
    os.environ["AWS_SECRET_ACCESS_KEY"] = "minioadmin"
    os.environ["AWS_ENDPOINT_URL"] = "http://localhost:9000"

    # Create bucket first
    bucket_path = UPath("s3://docling-test")
    try:
        bucket_path.mkdir(exist_ok=True)
        print(f"  Created bucket: {bucket_path}")
    except Exception as e:
        print(f"  Bucket creation: {e}")

    output_dir = UPath("s3://docling-test/output")
    json_path = output_dir / "doc.json"
    artifacts_dir = output_dir / "images"

    try:
        doc.save_as_json(
            json_path,
            artifacts_dir=artifacts_dir,
            image_mode=ImageRefMode.REFERENCED
        )
        print(f"SUCCESS: Saved to {json_path}")
        print(f"  Images dir: {artifacts_dir}")

        # List saved files
        print("  Files in S3:")
        for f in output_dir.glob("**/*"):
            if f.is_file():
                print(f"    {f}")
    except TypeError as e:
        print(f"TYPE ERROR (expected if Path type check fails): {e}")
    except Exception as e:
        print(f"FAILED: {type(e).__name__}: {e}")
    print()


def test_upath_s3_roundtrip(doc: DoclingDocument):
    """Test saving and loading document from S3."""
    print("=== Test: UPath S3 Round-trip ===")

    # Configure for local MinIO
    os.environ["AWS_ACCESS_KEY_ID"] = "minioadmin"
    os.environ["AWS_SECRET_ACCESS_KEY"] = "minioadmin"
    os.environ["AWS_ENDPOINT_URL"] = "http://localhost:9000"

    output_dir = UPath("s3://docling-test/roundtrip")
    json_path = output_dir / "doc.json"
    yaml_path = output_dir / "doc.yaml"

    try:
        # Save as JSON
        doc.save_as_json(json_path, image_mode=ImageRefMode.EMBEDDED)
        print(f"  Saved JSON to: {json_path}")

        # Save as YAML
        doc.save_as_yaml(yaml_path, image_mode=ImageRefMode.EMBEDDED)
        print(f"  Saved YAML to: {yaml_path}")

        # Load back from JSON
        loaded_json = DoclingDocument.load_from_json(json_path)
        print(f"  Loaded from JSON: {loaded_json.name}")
        print(f"    Pages: {len(loaded_json.pages)}, Pictures: {len(loaded_json.pictures)}")

        # Load back from YAML
        loaded_yaml = DoclingDocument.load_from_yaml(yaml_path)
        print(f"  Loaded from YAML: {loaded_yaml.name}")
        print(f"    Pages: {len(loaded_yaml.pages)}, Pictures: {len(loaded_yaml.pictures)}")

        # Verify content matches
        assert len(loaded_json.pages) == len(doc.pages), "Page count mismatch (JSON)"
        assert len(loaded_yaml.pages) == len(doc.pages), "Page count mismatch (YAML)"
        assert len(loaded_json.pictures) == len(doc.pictures), "Picture count mismatch (JSON)"
        assert len(loaded_yaml.pictures) == len(doc.pictures), "Picture count mismatch (YAML)"

        print("SUCCESS: Round-trip verified!")
    except Exception as e:
        print(f"FAILED: {type(e).__name__}: {e}")
        import traceback
        traceback.print_exc()
    print()


def test_manual_upath_save(doc: DoclingDocument):
    """If native doesn't work, test manual save with UPath."""
    print("=== Test: Manual UPath S3 Save ===")

    os.environ["AWS_ACCESS_KEY_ID"] = "minioadmin"
    os.environ["AWS_SECRET_ACCESS_KEY"] = "minioadmin"
    os.environ["AWS_ENDPOINT_URL"] = "http://localhost:9000"

    output_dir = UPath("s3://docling-test/manual")

    try:
        output_dir.mkdir(parents=True, exist_ok=True)

        # Save JSON locally first, then upload
        import json
        json_content = doc.model_dump_json(indent=2)
        json_path = output_dir / "doc.json"
        json_path.write_text(json_content)
        print(f"  Saved JSON to: {json_path}")

        # Save images manually
        for i, picture in enumerate(doc.pictures[:3]):  # First 3
            if picture.image and picture.image.pil_image:
                from io import BytesIO
                buf = BytesIO()
                picture.image.pil_image.save(buf, format="PNG")
                img_path = output_dir / f"picture_{i}.png"
                img_path.write_bytes(buf.getvalue())
                print(f"  Saved image: {img_path}")

        print("SUCCESS: Manual save worked")
    except Exception as e:
        print(f"FAILED: {type(e).__name__}: {e}")
    print()


if __name__ == "__main__":
    test_type_compatibility()

    print("Converting document...")
    doc = convert_doc()
    print(f"Document has {len(doc.pages)} pages, {len(doc.pictures)} pictures\n")

    test_local_path(doc)
    test_upath_local(doc)
    test_upath_s3(doc)
    test_upath_s3_roundtrip(doc)
    test_manual_upath_save(doc)